from torch.utils.data import Dataset
from PIL import Image
import os


class MyDataset(Dataset):

    def __init__(self, root_dir, label_dir):
        self.root_dir = root_dir
        self.label_dir = label_dir
        self.path = os.path.join(self.root_dir, self.label_dir)
        self.images = os.listdir(self.path)

    def __getitem__(self, index):
        img_path = os.path.join(self.path, self.images[index])
        img = Image.open(img_path)
        return img, self.label_dir

    def __len__(self):
        return len(self.images)


root_dir = "./dataset/train"
ants_label = "ants"
bees_label = "bees"

ants_dataset = MyDataset(root_dir, ants_label)
bees_dataset = MyDataset(root_dir, bees_label)
print("all_dataset:%s" % len(ants_dataset))
print("bees_dataset:%s" % len(bees_dataset))

all_dataset = ants_dataset + bees_dataset
print("all_dataset:%s" % len(all_dataset))

img, label = all_dataset[125]
img.show()
print(label)